from re import X
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
import torch
from Models.models.randaugment import RandAugmentMC
# from timm.data.auto_augment import rand_augment_transform #可使用rand_augment应该有同样效果

class TransformFixMatch(object):
    def __init__(self, image_size=84):
        self.weak = transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop(size=image_size,
                                  padding=int(image_size*0.125),
                                  padding_mode='reflect')])
        self.strong = transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop(size=image_size,
                                  padding=int(image_size*0.125),
                                  padding_mode='reflect'),
            RandAugmentMC(n=2, m=10)])
        self.normalize = transforms.Compose([
            transforms.ToTensor(),
            # normalize 在数据读入时已经完成
            # transforms.Normalize(mean=mean, std=std)
            ])

    def __call__(self, x):
        if isinstance(x, torch.Tensor):
            x = transforms.ToPILImage()(x)
        weak = self.weak(x)
        strong = self.strong(x)
        return self.normalize(weak), self.normalize(strong)

class RandAug(object):
    def __init__(self, image_size=84):
        self.strong = transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop(size=image_size,
                                  padding=int(image_size*0.125),
                                  padding_mode='reflect'),
            RandAugmentMC(n=2, m=10)])
        self.normalize = transforms.Compose([
            transforms.ToTensor(),
            # normalize 在数据读入时已经完成
            # transforms.Normalize(mean=mean, std=std)
            ])

    def __call__(self, x):
        if isinstance(x, torch.Tensor):
            x = transforms.ToPILImage()(x)
        strong = self.strong(x)
        return self.normalize(strong)